﻿using System;
using System.Collections.Generic;
using Allegro.Framework;
using Allegro.Integrator;
//using Cosmologica.Framework;

namespace Allegro.Integrator
{
    public class RungeKuttaIntegrator : IODEIntegrator
    {
        public const double TINY = 1.0e-30;

        /// <summary>
        /// Default accuracy
        /// </summary>
        public const double DefaultAccuracy = 1.0e-8;
        const double SAFETY = 0.9;
        const double PGROW = -0.2;
        const double PSHRNK = -0.25;
        const double ERRCON = 1.89e-4; // The value ERRCON equals (5/SAFETY) raised to the power (1/PGROW), see use below
        const double MaxStepReductionFactor = 0.1;
        const double MaxStepIncreaseFactor = 5;

        static readonly double a2 = 0.2, a3 = 0.3, a4 = 0.6, a5 = 1.0, a6 = 0.875;
        static readonly double b21 = 0.2, b31 = 3.0 / 40.0, b32 = 9.0 / 40.0;
        static readonly double b41 = 0.3, b42 = -0.9, b43 = 1.2;
        static readonly double b51 = -11.0 / 54.0, b52 = 2.5;
        static readonly double b53 = -70.0 / 27.0, b54 = 35.0 / 27.0;
        static readonly double b61 = 1631.0 / 55296.0, b62 = 175.0 / 512.0, b63 = 575.0 / 13824.0;
        static readonly double b64 = 44275.0 / 110592.0, b65 = 253.0 / 4096.0;
        static readonly double c1 = 37.0 / 378.0, c3 = 250.0 / 621.0, c4 = 125.0 / 594.0, c6 = 512.0 / 1771.0;
        static readonly double dc5 = -277.0 / 14336.0;
        static readonly double dc1 = c1 - 2825.0 / 27648.0, dc3 = c3 - 18575.0 / 48384.0;
        static readonly double dc4 = c4 - 13525.0 / 55296.0, dc6 = c6 - 0.25;

        int _n;
        double[] _ak2;
        double[] _ak3;
        double[] _ak4;
        double[] _ak5;
        double[] _ak6;
        double[] _ytemp;
        double[] _yComputed;
        double[] _yerr;
        double[] _dydx;
        double[] _yscal;

        double _accuracy = DefaultAccuracy;

        bool[] _integrated; // Flag for each parameter. If true, variable i was integrated.
        bool[] _computed; // Flag for each parameter. If true, variable i was computed directly

        /// <summary>
        /// 
        /// </summary>
        /// <param name="dimension"></param>
        public RungeKuttaIntegrator(int dimension)
        {
            if (dimension <= 0)
                throw new ApplicationException("Dim. must be > 0");

            _n = dimension;
            _ak2 = new double[_n];
            _ak3 = new double[_n];
            _ak4 = new double[_n];
            _ak5 = new double[_n];
            _ak6 = new double[_n];
            _ytemp = new double[_n];
            _yComputed = new double[_n];
            _yerr = new double[_n];
            _dydx = new double[_n];
            _yscal = new double[_n];

            _integrated = new bool[_n];
            _computed = new bool[_n];
            _ResetCompletedFlags();
        }

        /// <summary>
        /// Relative accuracy
        /// </summary>
        public double Accuracy
        {
            get
            {
                return _accuracy;
            }
            set
            {
                _accuracy = value;
            }
        }

        /// <summary>
        /// Integration step. Input are the dependent variable vector y[1..n] and its derivative dydx[1..n]
        /// at the starting value of the independent varable x. Also input are the stepsize to be attempted, 
        /// htry, and the vector yscal[1..n], against which the error is scaled. On output, y and x are replaced by their
        /// new values, hdid is the stepsize that was actually accomplished, and hnext is the estimated next stepsize. 
        /// derivs is the user-supplied method that computes the right-hand side derivatives
        /// </summary>        
        /// <param name="x"></param>
        /// <param name="y"></param>
        /// <param name="htry"></param>
        /// <param name="hdid"></param>
        /// <param name="hnext"></param>
        /// <param name="derivs"></param>
        /// <param name="values"></param>
        public void Integrate(ref double x, ref double[] y, double htry, out double hdid, out double hnext, 
            Derivative derivs, ComputeValues values)
        {
            double errmax, h, htemp, xnew;
            double eps = Accuracy;

            _ResetCompletedFlags();
            bool isComplete = derivs(x, y, ref _dydx, ref _integrated);

            h = htry; // Set stepsize to initial trial value
            for (int i = 0; i < _n; i++)
            {
                // Scaling used to monitor accuracy. This general-purpose choice can be modified if need be.
                _yscal[i] = Math.Abs(y[i]) + Math.Abs(_dydx[i] * h) + TINY;
            }

            //double[] errors = new double[_n];
            //int iErrorMax = 0;
            for (; ; )
            {
                // Take a step
                _MakeIntegrationStep(y, _dydx, x, h, ref _ytemp, ref _yerr, derivs, values);
                errmax = 0.0;
                for (int i = 0; i < _n; i++)
                {
                    //errors[i] = Math.Abs(_yerr[i] / _yscal[i])/eps;
                    //errors[i] = Math.Abs(_yerr[i]) / eps;
                    //if (errors[i] > errmax / eps)
                    //    iErrorMax = i;
                    //errmax = Math.Max(errmax, Math.Abs(_yerr[i] / _yscal[i]));
                    errmax = Math.Max(errmax, Math.Abs(_yerr[i]));
                }
                errmax /= eps;
                if (errmax <= 1.0)
                    break;
                htemp = SAFETY * h * Math.Pow(errmax, PSHRNK);
                // Truncation error too large, reduce stepsize
                // Limit the step reduction by MaxStepReductionFactor (0.1 by default)
                h = (h > 0.0 ? Math.Max(htemp, MaxStepReductionFactor * h) : Math.Min(htemp, MaxStepReductionFactor * h));
                xnew = x + h;
                if (xnew == x)
                    ErrorHandler.Error(string.Format("Stepsize underflow in RungeKuttaIntegrator.Integrate occurred at x={0}", x));
            }
            if (errmax > ERRCON)
                hnext = SAFETY * h * Math.Pow(errmax, PGROW);
            else
                hnext = MaxStepIncreaseFactor * h; // No more than a factor of 5 increase
            x += h;
            hdid = h;
            for (int i = 0; i < _n; i++)
            {
                y[i] = _ytemp[i];
            }
        }

        // Private methods
        void _ResetCompletedFlags()
        {
            for (int i = 0; i < _n; i++)
            {
                _computed[i] = _integrated[i] = false;
            }
        }
        /// <summary>
        /// Given values for n variables y[1..n] and theri derivatives dydx[1..n] known at x, use 
        /// the fifth-order Cash-Karp Runge-Kutta method to advance the solution over an interval h and
        /// return the incremented variables as yout[1..n]. Also return an estimate of the local
        /// truncation error in yout using the embedded fourth-order method. The method invokes the 
        /// user-supplied method derivs through the delegat Derivative(x,y,dydx), 
        /// which computes the first-order derivatives at x.
        /// </summary>
        /// <param name="y"></param>
        /// <param name="dydx"></param>
        /// <param name="x"></param>
        /// <param name="h"></param>
        /// <param name="yout"></param>
        /// <param name="yerr"></param>
        /// <param name="derivs"></param>
        void _MakeIntegrationStep(double[] y, double[] dydx, double x, double h, ref double[] yout, ref double[] yerr, Derivative derivs, ComputeValues values)
        {
            double xtemp;
            bool isComplete = false;
            // First step
            for (int i = 0; i < _n; i++) 
            {
                if (_integrated[i])
                    _ytemp[i] = y[i] + h * b21 * dydx[i];
            }
            if (!isComplete && values != null)
                isComplete = values(x, y, ref _yComputed, _integrated, ref _computed);
            for (int i = 0; i < _n; i++)
            {
                if (!_integrated[i] && _computed[i])
                    _ytemp[i] = _yComputed[i];
            }
            // Second step
            _ResetCompletedFlags();
            xtemp = x + a2 * h;
            isComplete = derivs(xtemp, _ytemp, ref _ak2, ref _integrated);
            for (int i = 0; i < _n; i++)
            {
                if (_integrated[i])
                    _ytemp[i] = y[i] + h * (b31 * dydx[i] + b32 * _ak2[i]);
            }
            if (!isComplete && values != null)
                isComplete = values(xtemp, _ytemp, ref _yComputed, _integrated, ref _computed);
            for (int i = 0; i < _n; i++)
            {
                if (!_integrated[i] && _computed[i])
                    _ytemp[i] = _yComputed[i];
            }
            // Third step
            _ResetCompletedFlags();
            xtemp = x + a3 * h;
            isComplete = derivs(xtemp, _ytemp, ref _ak3, ref _integrated);
            for (int i = 0; i < _n; i++)
            {
                if (_integrated[i])
                    _ytemp[i] = y[i] + h * (b41 * dydx[i] + b42 * _ak2[i] + b43 * _ak3[i]);
            }
            if (!isComplete && values != null)
                isComplete = values(xtemp, _ytemp, ref _yComputed, _integrated, ref _computed);
            for (int i = 0; i < _n; i++)
            {
                if (!_integrated[i] && _computed[i])
                    _ytemp[i] = _yComputed[i];
            }
            // Fourth step
            _ResetCompletedFlags();
            xtemp = x + a4 * h;
            isComplete = derivs(xtemp, _ytemp, ref _ak4, ref _integrated);
            for (int i = 0; i < _n; i++)
            {
                if (_integrated[i])
                    _ytemp[i] = y[i] + h * (b51 * dydx[i] + b52 * _ak2[i] + b53 * _ak3[i] + b54 * _ak4[i]);
            }
            if (!isComplete && values != null)
                isComplete = values(xtemp, _ytemp, ref _yComputed, _integrated, ref _computed);
            for (int i = 0; i < _n; i++)
            {
                if (!_integrated[i] && _computed[i])
                    _ytemp[i] = _yComputed[i];
            }
            // Fifth step
            _ResetCompletedFlags();
            xtemp = x + a5 * h;
            isComplete = derivs(xtemp, _ytemp, ref _ak5, ref _integrated);
            for (int i = 0; i < _n; i++)
            {
                if (_integrated[i])
                    _ytemp[i] = y[i] + h * (b61 * dydx[i] + b62 * _ak2[i] + b63 * _ak3[i] + b64 * _ak4[i] + b65 * _ak5[i]);
            }
            if (!isComplete && values != null)
                isComplete = values(xtemp, _ytemp, ref _yComputed, _integrated, ref _computed);
            for (int i = 0; i < _n; i++)
            {
                if (!_integrated[i] && _computed[i])
                    _ytemp[i] = _yComputed[i];
            }
            // Sixth step
            _ResetCompletedFlags();
            xtemp = x + a6 * h;
            isComplete = derivs(xtemp, _ytemp, ref _ak6, ref _integrated);
            // Accumulate increments with proper weights
            for (int i = 0; i < _n; i++)
            {
                if (_integrated[i])
                    yout[i] = y[i] + h * (c1 * dydx[i] + c3 * _ak3[i] + c4 * _ak4[i] + c6 * _ak6[i]);
            }
            if (!isComplete && values != null)
                isComplete = values(xtemp, _ytemp, ref yout, _integrated, ref _computed);
            
            // Estimate error as difference between fourth and fifth order methods
            for (int i = 0; i < _n; i++)
            {
                if (_integrated[i])
                {
                    yerr[i] = h*(dc1*dydx[i] + dc3*_ak3[i] + dc4*_ak4[i] + dc5*_ak5[i] + dc6*_ak6[i])/_yscal[i];
                }
                else
                    yerr[i] = 0.0;
            }
        }
    }
}
